import torch
print("test cuda")
if torch.cuda.is_available():
    print("CUDA is available! Training on GPU...")
else:
    print("CUDA is not available. Training on CPU...")